[Main] Support device-init grouped linear module too with when not using TE opfuser by zhongbozhu · Pull Request #6000 · NVIDIA/Megatron-LM

zhongbozhu · 2026-07-23T20:54:48Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Proper single weight support, discrete weight is not affected. Adds many numerical unit tests running on GB200.

Unit testing:

Note: some tests can only run with TE updated (NVIDIA/TransformerEngine#3224), for CI image, it should just pass without causing errors.

pytest -s -v --confcutdir=tests/unit_tests/transformer/moe tests/unit_tests/transformer/moe/test_grouped_mlp.py

torchrun --nproc_per_node=4 --log-dir /tmp/mcore-single-weight-ut --tee 0:3 --redirects 3 -m pytest -s -v --confcutdir=tests/unit_tests/transformer/moe tests/unit_tests/transformer/moe/test_moe_single_grouped_weight_numerics.py

torchrun --nproc_per_node=4 --log-dir /tmp/mcore-grouped-dispatcher-ut --tee 0:3 --redirects 3 -m pytest -s -v --confcutdir=tests/unit_tests/transformer/moe tests/unit_tests/transformer/moe/test_grouped_tensor_dispatcher_numerics.py

E2E testing:

Model - Qwen3.5 VL SFT
Dispatcher - HybridEP
Experiements - bf16 cublas grouped gemm [single weight ON/OFF] vs. cuteDSL with TE op fuser mxfp8 [single weight ON/OFF] vs mxfp8 cublas grouped gemm [single weight ON/OFF]

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

copy-pr-bot · 2026-07-23T20:54:52Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

zhongbozhu · 2026-08-01T09:13:21Z

/ok to test 1f647e6

zhongbozhu · 2026-08-03T21:16:03Z

            )

+    @staticmethod
+    def _apply_packed_bias(intermediate_parallel, packed_bias, tokens_per_expert, permuted_probs):


Note: only models with routed expert bias will use this

zhongbozhu · 2026-08-04T01:50:24Z

/ok to test d5c9aed

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

zhongbozhu · 2026-08-05T18:40:50Z

/ok to test f921404

YangFei1990

Just for my understanding, mcore will pass moe_use_grouped_tensor / use_transformer_engine_op_fuser into TE, but it is TE's responsibility to pick the groupgemm backend, which is depending on hardware/dtype/envs, so it might not necessarily follow what exactly provided by mcore, is that right?

YangFei1990 · 2026-08-05T21:12:34Z

+        # Some dispatchers have already padded each expert's token segment before the tokens
+        # reach this module:
+        #   * router padding changes the routing map before dispatch;
+        #   * HybridEP pads as part of its fused dispatch/permute operation;


nit: HybridEP/NCCL EP

YangFei1990 · 2026-08-05T21:14:29Z

+        # The token buffer may already contain per-expert padding when padding was performed
+        # before expert compute:
+        #   * router padding modified the routing map before dispatch;
+        #   * HybridEP fused padding into dispatch/permute;


nit: same + NCCL EP

YangFei1990 · 2026-08-05T21:16:02Z

+    assert torch.isfinite(hidden_states.grad).all()
+
+
+class TestGroupedTensorDispatcherNumerics:


Could we also add test cases for NCCL EP?

yaox12 approved these changes Jul 27, 2026

View reviewed changes

zhongbozhu mentioned this pull request Jul 29, 2026

GroupedLinear(single_grouped_weight=True) produces no weight gradient (wgrad) NVIDIA/TransformerEngine#3234

Open

zhongbozhu force-pushed the main_support_device_group_linear branch 2 times, most recently from 7294e71 to 1f4a49e Compare July 29, 2026 02:32

zhongbozhu mentioned this pull request Jul 31, 2026

Improve device-init grouped linear module with single grouped weight support NVIDIA/TransformerEngine#3224

Open

13 tasks

zhongbozhu force-pushed the main_support_device_group_linear branch from 944cb4e to a78c12b Compare August 1, 2026 00:55

zhongbozhu marked this pull request as ready for review August 1, 2026 09:12

zhongbozhu requested review from a team as code owners August 1, 2026 09:12

copy-pr-bot Bot temporarily deployed to public August 1, 2026 09:14 Inactive

copy-pr-bot Bot temporarily deployed to test August 1, 2026 09:14 Inactive

github-actions Bot added the complexity: medium label Aug 1, 2026

copy-pr-bot Bot temporarily deployed to public August 1, 2026 09:17 Inactive

copy-pr-bot Bot temporarily deployed to public August 1, 2026 09:29 Inactive

zhongbozhu commented Aug 3, 2026

View reviewed changes

copy-pr-bot Bot temporarily deployed to public August 4, 2026 01:51 Inactive

copy-pr-bot Bot temporarily deployed to test August 4, 2026 01:51 Inactive

copy-pr-bot Bot temporarily deployed to public August 4, 2026 01:54 Inactive

copy-pr-bot Bot temporarily deployed to public August 4, 2026 02:05 Inactive

zhongbozhu requested review from a team as code owners August 5, 2026 07:09

zhongbozhu force-pushed the main_support_device_group_linear branch from 57ff027 to d6d76b2 Compare August 5, 2026 07:10

zhongbozhu added 7 commits August 5, 2026 11:14

support device group linear

a3cdc25

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

improve tests

b239e37

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

fix unit test failure

bf6095a

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

pass all tests

225be1c

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

format

87ca72a

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

call te_general_gemm properly

e7f7996

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

fix CI issues

f921404

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

zhongbozhu force-pushed the main_support_device_group_linear branch from d6d76b2 to f921404 Compare August 5, 2026 18:40

copy-pr-bot Bot temporarily deployed to public August 5, 2026 18:41 Inactive

copy-pr-bot Bot temporarily deployed to test August 5, 2026 18:41 Inactive

copy-pr-bot Bot temporarily deployed to public August 5, 2026 18:44 Inactive

copy-pr-bot Bot temporarily deployed to public August 5, 2026 18:45 Inactive

copy-pr-bot Bot temporarily deployed to public August 5, 2026 18:54 Inactive

YangFei1990 reviewed Aug 5, 2026

View reviewed changes

Navigation Menu

Uh oh!

[Main] Support device-init grouped linear module too with when not using TE opfuser - #6000

[Main] Support device-init grouped linear module too with when not using TE opfuser#6000
zhongbozhu wants to merge 7 commits into
NVIDIA:mainfrom
zhongbozhu:main_support_device_group_linear

zhongbozhu commented Jul 23, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Jul 23, 2026

Uh oh!

zhongbozhu commented Aug 1, 2026

Uh oh!

zhongbozhu Aug 3, 2026

Uh oh!

zhongbozhu commented Aug 4, 2026

Uh oh!

zhongbozhu commented Aug 5, 2026

Uh oh!

YangFei1990 left a comment

Uh oh!

YangFei1990 Aug 5, 2026

Uh oh!

YangFei1990 Aug 5, 2026

Uh oh!

YangFei1990 Aug 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		assert torch.isfinite(hidden_states.grad).all()


		class TestGroupedTensorDispatcherNumerics:

Uh oh!

Conversation

zhongbozhu commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Jul 23, 2026

Uh oh!

zhongbozhu commented Aug 1, 2026

Uh oh!

zhongbozhu Aug 3, 2026

Choose a reason for hiding this comment

Uh oh!

zhongbozhu commented Aug 4, 2026

Uh oh!

zhongbozhu commented Aug 5, 2026

Uh oh!

YangFei1990 left a comment

Choose a reason for hiding this comment

Uh oh!

YangFei1990 Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

YangFei1990 Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

YangFei1990 Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

zhongbozhu commented Jul 23, 2026 •

edited

Loading